# Arguments:
# $args[0] = ScheduleDay
# $args[1] = ScheduleTime
# $args[2] = ScanType
# $args[3] = ScanOnlyWhenIdle
# $args[4] = ScanLocation (optional)

if ($args.Count -lt 4) {
    Write-Output "ERROR: Usage: <ScheduleDay> <ScheduleTime> <ScanType> <ScanOnlyWhenIdle> [ScanLocation]"
    exit 1
}

$ScheduleDay       = $args[0]
$ScheduleTime      = $args[1]
$ScanType          = $args[2]
$ScanOnlyWhenIdle  = $args[3]
$ScanLocation      = if ($args.Count -ge 5) { $args[4] } else { "" }

try {
    # Validate values
    $validDays = @("Everyday","Sunday","Monday","Tuesday","Wednesday","Thursday","Friday","Saturday","Never")
    if ($ScheduleDay -notin $validDays) {
        throw "Invalid ScheduleDay: $ScheduleDay"
    }

    $validTypes = @("Quick","Full","Custom")
    if ($ScanType -notin $validTypes) {
        throw "Invalid ScanType: $ScanType"
    }

    if ($ScanOnlyWhenIdle -notin @("true","false")) {
        throw "ScanOnlyWhenIdle must be 'true' or 'false'."
    }

    # Validate ScanLocation for Custom scan
    if ($ScanType -eq "Custom") {
        if ([string]::IsNullOrWhiteSpace($ScanLocation)) {
            throw "ScanLocation is required when ScanType is 'Custom'."
        }

        # Layer 1 — Whitelist: only allow safe path characters (A03 Injection mitigation)
        if ($ScanLocation -notmatch '^[A-Za-z]:\\[\w\\ .\-]+$') {
            throw "ScanLocation contains invalid characters. Only alphanumeric paths are permitted."
        }

        if (-not (Test-Path -Path $ScanLocation -ErrorAction Stop)) {
            throw "ScanLocation '$ScanLocation' does not exist."
        }
    }

    # Check Windows Defender service is running
    $defenderService = Get-Service -Name WinDefend -ErrorAction SilentlyContinue
    if (-not $defenderService -or $defenderService.Status -ne 'Running') {
        throw "Windows Defender service (WinDefend) is not running. Ensure Microsoft Defender Antivirus is enabled and the service is active."
    }

    # Parse time — strict HH:mm format only; rejects date-embedded strings and ambiguous formats
    $parsedTime = $null
    if (-not [datetime]::TryParseExact(
            $ScheduleTime,
            "HH:mm",
            [System.Globalization.CultureInfo]::InvariantCulture,
            [System.Globalization.DateTimeStyles]::None,
            [ref]$parsedTime)) {
        throw "Invalid ScheduleTime '$ScheduleTime'. Expected format: HH:mm (e.g. 14:30)."
    }

    # Convert idle flag
    $idleBool = ($ScanOnlyWhenIdle.ToLower() -eq "true")

    if ($ScanType -eq "Quick" -or $ScanType -eq "Full") {

        $dayMap = @{
            "Everyday"  = 0
            "Sunday"    = 1
            "Monday"    = 2
            "Tuesday"   = 3
            "Wednesday" = 4
            "Thursday"  = 5
            "Friday"    = 6
            "Saturday"  = 7
            "Never"     = 8
        }

        $scheduleDayValue = $dayMap[$ScheduleDay]
        $scanTypeValue = @{ "Quick" = 1; "Full" = 2 }[$ScanType]

        Set-MpPreference `
            -ScanScheduleDay $scheduleDayValue `
            -ScanScheduleTime $parsedTime `
            -ScanParameters $scanTypeValue `
            -ScanOnlyIfIdleEnabled $idleBool `
            -ErrorAction Stop

        Write-Output "Scheduled scan configured: Day=$ScheduleDay | Time=$ScheduleTime | Type=$ScanType | IdleOnly=$ScanOnlyWhenIdle"
    }
    elseif ($ScanType -eq "Custom") {

        $taskName = "WindowsDefender_CustomScan"
        $normalizedPath = $ScanLocation.TrimEnd('\','/')
        # Layer 2 — escape any residual single quotes so they cannot break the command boundary
        $escapedPath = $normalizedPath -replace "'", "''"

        if ($ScheduleDay -eq "Never") {

            $existing = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue

            if ($existing) {
                Disable-ScheduledTask -TaskName $taskName -ErrorAction Stop
                Write-Output "Custom location scan task disabled."
            }
            else {
                Write-Output "No existing custom location scan task found to disable."
            }
        }
        else {

            # Layer 3 — use -EncodedCommand so no shell metacharacter in the path can be misinterpreted
            $cmdBytes   = [System.Text.Encoding]::Unicode.GetBytes("Start-MpScan -ScanType CustomScan -ScanPath '$escapedPath'")
            $encodedCmd = [Convert]::ToBase64String($cmdBytes)

            $action = New-ScheduledTaskAction `
                -Execute "powershell.exe" `
                -Argument "-NonInteractive -WindowStyle Hidden -EncodedCommand $encodedCmd"

            if ($ScheduleDay -eq "Everyday") {
                $trigger = New-ScheduledTaskTrigger -Daily -At $parsedTime
            }
            else {
                $trigger = New-ScheduledTaskTrigger -Weekly -DaysOfWeek $ScheduleDay -At $parsedTime
            }

            $settings = New-ScheduledTaskSettingsSet `
                -RunOnlyIfIdle:$idleBool `
                -IdleDuration (New-TimeSpan -Minutes 10)

            $principal = New-ScheduledTaskPrincipal `
                -UserId "SYSTEM" `
                -LogonType ServiceAccount `
                -RunLevel Highest

            Register-ScheduledTask `
                -TaskName $taskName `
                -Action $action `
                -Trigger $trigger `
                -Settings $settings `
                -Principal $principal `
                -Force `
                -ErrorAction Stop

            Write-Output "Custom location scan scheduled: Day=$ScheduleDay | Time=$ScheduleTime | Location=$normalizedPath | IdleOnly=$ScanOnlyWhenIdle"
        }
    }
}
catch {
    Write-Output "ERROR: $($_.Exception.Message)"
    exit 1
}